from torch.utils.data import Dataset
import torch
import cv2
import os
from typing import Optional,List,Tuple,Dict
import torchvision
import numpy as np

CLASS_NAMES:Dict[str,int]={
    'ADI':0,
    'BACK':1,
    'DEB':2,
    'LYM':3,
    'MUC':4,
    'MUS':5,
    'NORM':6,
    'STR':7,
    'TUM':8,
    }

class cancerClassDataset(Dataset):
    def __init__(self, root_folder:str, \
            raw_transform:Optional[torchvision.transforms.transforms.Compose]=None) -> None:
        
        
        self.raw_transform=raw_transform
        if self.raw_transform == None:
            self.raw_transform=torchvision.transforms.ToTensor()

        self.all_pairs:List[Tuple[int]]=[]
        for class_name, ID in CLASS_NAMES.items():
            class_name:str
            ID:int
            tif_files:List[str]= os.listdir(os.path.join(root_folder,class_name))
            for tif_file in tif_files:
                tif_file:str
                if os.path.isfile(os.path.join(root_folder,class_name,tif_file))==False\
                    or (tif_file.endswith('.tif')==False and tif_file.endswith('.tiff')==False):
                    continue
                assert(tif_file.startswith(class_name))
                self.all_pairs.append((os.path.join(root_folder,class_name,tif_file),ID,))

        super().__init__()

    def __getitem__(self, index: int) -> torch.Tensor:
        raw_img:np.ndarray=cv2.imread(self.all_pairs[index][0],cv2.IMREAD_UNCHANGED)
        label:int=self.all_pairs[index][1]
        return self.raw_transform(raw_img), label

    def __len__(self)->int:
        return len(self.all_pairs)

if __name__ == '__main__':
    test=cancerClassDataset('E:/JohnsonProj/smallTraining')
    print(test[1700][0].shape,test[1700][1])
    print(len(test))